# agent.py - 交互式 AI 编程助手（支持多轮任务 & 后台命令）
import os
import json
import subprocess
from dotenv import load_dotenv
import dashscope  # 阿里云通义千问 SDK

# 加载环境变量（推荐在 .env 中设置 DASHSCOPE_API_KEY）
load_dotenv()

# 设置你的 DashScope API Key
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
if not dashscope.api_key:
    raise ValueError("请设置环境变量 DASHSCOPE_API_KEY")

# 工作空间目录
WORKSPACE = "workspace"

# 确保 workspace 目录存在
os.makedirs(WORKSPACE, exist_ok=True)

# === 全局字典：保存后台进程，便于管理 ===
background_processes = {}

# === 工具函数：AI 的“手脚” ===

def write_file(path: str, content: str) -> str:
    """写文件到 workspace"""
    full_path = os.path.join(WORKSPACE, path)
    try:
        with open(full_path, 'w', encoding='utf-8') as f:
            f.write(content)
        return f"✅ 成功写入文件: {path}"
    except Exception as e:
        return f"❌ 写入失败 {path}: {str(e)}"

def read_file(path: str) -> str:
    """读取 workspace 中的文件"""
    full_path = os.path.join(WORKSPACE, path)
    try:
        with open(full_path, 'r', encoding='utf-8') as f:
            content = f.read()
        return f"📄 文件 {path} 内容:\n{content}"
    except FileNotFoundError:
        return f"❌ 文件不存在: {path}"
    except Exception as e:
        return f"❌ 读取失败: {str(e)}"

def run_command(command: str, background: bool = False) -> str:
    """
    在 workspace 目录下执行命令
    - background=False: 前台阻塞执行，带超时控制（默认）
    - background=True: 使用 Popen 启动后台非阻塞进程
    """
    try:
        if background:
            proc = subprocess.Popen(
                command,
                shell=True,
                cwd=WORKSPACE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,
                universal_newlines=True
            )
            pid = proc.pid
            background_processes[pid] = proc
            return f"🟢 已在后台启动命令 (PID {pid}):\n命令: {command}\n💡 提示：该进程将持续运行，可通过 KILL_PROCESS 终止。"
        else:
            result = subprocess.run(
                command,
                shell=True,
                cwd=WORKSPACE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=30  # 多数脚本应在30秒内完成
            )
            if result.returncode == 0:
                return f"🟢 命令成功执行:\n{result.stdout.strip()}"
            else:
                return f"🔴 命令执行失败 (退出码 {result.returncode}):\n{result.stderr.strip()}"

    except subprocess.TimeoutExpired:
        return "⏰ 命令执行超时（>30秒），已终止。"
    except Exception as e:
        return f"💥 执行异常: {str(e)}"

def list_background_processes() -> str:
    """列出当前所有后台运行的进程"""
    alive_pids = []
    for pid, proc in list(background_processes.items()):
        if proc.poll() is None:
            alive_pids.append(f"PID {pid}: 运行中")
        else:
            del background_processes[pid]
    if alive_pids:
        return "🔍 当前后台进程:\n" + "\n".join(alive_pids)
    else:
        return "📭 没有正在运行的后台进程。"

def kill_process(pid_str: str) -> str:
    """终止指定 PID 的后台进程"""
    try:
        pid = int(pid_str)
        proc = background_processes.get(pid)
        if not proc:
            return f"❌ PID {pid} 不在监控列表中。"
        if proc.poll() is None:
            proc.terminate()  # 先尝试优雅退出
            try:
                proc.wait(timeout=3)
                del background_processes[pid]
                return f"✅ 成功终止进程 PID {pid}"
            except subprocess.TimeoutExpired:
                proc.kill()
                del background_processes[pid]
                return f"⚠️  强制杀死了进程 PID {pid}"
        else:
            del background_processes[pid]
            return f"ℹ️  PID {pid} 已自然结束。"
    except ValueError:
        return "❌ PID 必须是一个数字。"
    except Exception as e:
        return f"❌ 终止失败: {str(e)}"


# === 主控循环 ===

def ai_step(messages):
    """调用 Qwen-Plus 获取下一步动作"""
    try:
        response = dashscope.Generation.call(
            model="qwen-plus",
            messages=messages,
            temperature=0.3,
            max_tokens=4096,
            result_format='message'
        )
        if response.status_code == 200:
            return response.output.choices[0].message.content.strip()
        else:
            return f"❌ 调用 Qwen 失败: {response.message} (状态码: {response.status_code})"
    except Exception as e:
        return f"❌ 请求异常: {str(e)}"


def run_task_interactive(goal: str, full_history: list):
    """
    执行单个任务的主循环（共享全局对话历史）
    - full_history: 跨任务持久化的对话列表
    """
    print(f"\n{'='*60}")
    print(f"🎯 开始新任务: {goal}")
    print(f"{'='*60}")

    # 如果还没有系统提示，则添加一次
    if not full_history or full_history[0]["role"] != "system":
        system_prompt = {
            "role": "system",
            "content": '''
你是一个智能的 AI 编程助手，能够执行用户的编程任务并与环境交互。
你可以通过以下 JSON 指令操作：

{"action": "WRITE_FILE", "path": "xxx.py", "content": "..."}
{"action": "READ_FILE", "path": "xxx.py"}
{"action": "RUN_COMMAND", "command": "python xxx.py"}                    ← 前台运行
{"action": "RUN_COMMAND", "command": "python app.py", "background": true}  ← 后台运行服务
{"action": "LIST_PROCESSES"}
{"action": "KILL_PROCESS", "pid": 12345}
{"action": "THINK", "reason": "..."}
{"action": "FINAL_ANSWER", "answer": "..."}

📌 注意事项：
- 所有输出必须是纯 JSON 对象，不要包裹在代码块中，不要加额外解释。
- 若运行服务器等长期任务，请使用 `"background": true`。
- 文件均存于 workspace 目录下。
- 其中如果文件需要执行在控制台下，需要注意中文乱码的问题，我指的是代码中如果出现中文，然后需要采用 javac 编译的话，需要添加参数 javac -encoding utf-8

🧠 当前你具备记忆能力，可以看到此前的所有交互历史。
🎯 用户最新指令是你的最高优先级任务。
请立即开始规划并执行。
            '''.strip()
        }
        full_history.append(system_prompt)

    # 添加用户本次输入
    full_history.append({"role": "user", "content": goal})

    max_steps = 20
    task_success = False

    for step in range(max_steps):
        print(f"🔄 步骤 {step + 1}:")

        # 使用完整的对话历史请求 AI 决策
        ai_output = ai_step(full_history)
        print(f"🤖 AI 输出:\n{ai_output}")

        # 记录 AI 的响应（即使不是合法 JSON，也先记下来）
        full_history.append({"role": "assistant", "content": ai_output})

        # 尝试解析动作
        try:
            action = json.loads(ai_output)
        except json.JSONDecodeError:
            print("⚠️  AI 输出不是有效 JSON，尝试继续...\n")
            continue

        action_type = action.get("action")

        # === 执行对应动作 ===
        if action_type == "WRITE_FILE":
            path = action.get("path", "").strip()
            content = action.get("content", "")
            result = write_file(path, content) if path else "❌ 缺少 path 参数"

        elif action_type == "READ_FILE":
            path = action.get("path", "").strip()
            result = read_file(path)

        elif action_type == "RUN_COMMAND":
            cmd = action.get("command", "").strip()
            bg = action.get("background", False)
            result = run_command(cmd, background=bg) if cmd else "❌ 缺少 command 参数"
            if cmd:
                print(f"🔧 执行命令: {cmd} {'(后台)' if bg else '(前台)'}")

        elif action_type == "LIST_PROCESSES":
            result = list_background_processes()

        elif action_type == "KILL_PROCESS":
            pid = action.get("pid")
            result = kill_process(str(pid)) if pid is not None else "❌ 缺少 pid 参数"

        elif action_type == "THINK":
            reason = action.get("reason", "")
            result = f"[思考] {reason}"
            print(f"💭 {reason}")

        elif action_type == "FINAL_ANSWER":
            answer = action.get("answer", "")
            print(f"\n🎉 任务完成: {answer}")
            task_success = True
            break  # 成功完成任务

        else:
            result = f"❓ 未知动作类型: {action_type}"

        # 把执行结果反馈给 AI（加入历史，形成记忆）
        print(f"📊 执行结果:\n{result}\n")
        full_history.append({"role": "user", "content": f"执行结果:\n{result}"})

    if not task_success:
        print("⚠️  达到最大步骤数仍未完成任务。")

    return task_success


# === 启动入口（交互式）===
if __name__ == "__main__":
    print("🚀 欢迎使用交互式 AI 编程助手！")
    print("💡 提示：输入你的编程任务描述，例如：")
    print("     '写一个 Python 脚本，计算斐波那契数列前10项'")
    print("     '创建一个简单的网页并用 Python 启动服务器'")
    print("     输入 'exit' 或 'quit' 退出程序\n")

    # 清理上一次可能残留的进程引用
    background_processes.clear()

    # 🔁 全局对话历史：跨任务持久化记忆
    full_conversation_history = []

    while True:
        try:
            user_input = input("\n📝 请输入你的任务（或 exit 退出）: ").strip()

            if not user_input:
                print("⚠️  请输入有效任务内容。")
                continue

            if user_input.lower() in ['exit', 'quit', 'bye']:
                print("👋 再见！欢迎下次使用～")
                break

            if user_input.lower() in ['clear', 'cls']:
                os.system('cls' if os.name == 'nt' else 'clear')
                print("🚀 欢迎使用交互式 AI 编程助手！")
                print("💡 提示：输入你的编程任务描述，例如：")
                print("     '写一个 Python 脚本，计算斐波那契数列前10项'")
                print("     '创建一个简单的网页并用 Python 启动服务器'")
                print("     输入 'exit' 或 'quit' 退出程序\n")
                continue

            # ✅ 执行任务并传入/更新全局历史
            run_task_interactive(user_input, full_conversation_history)

        except KeyboardInterrupt:
            print("\n\n👋 检测到中断，正在退出...")
            break
        except Exception as e:
            print(f"❌ 程序发生异常: {e}")
